#include "ast.h"

#include <iostream>

#include "lexer.h"
#include "location.h"
#include "symbol.h"
#include "token.h"
#include "types.h"
#include "utils.h"

#define NAME BLUE
#define ARGUMENT GREY << "| " << DEF
#define VALUE GREEN

using Tag = Token::Tag;
using Utils::Indent;

namespace AST {
    void Expression::Print(int dep) const {
        Indent(dep), std::cout << NAME "Expression\n";
    }
    void Expression::PrintInfo(int dep) const {
        if (info) info->Print(dep + 1);
    }

    void AssignExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Assignment\n";
        PrintInfo(dep);
        Utils::Assert(lhs && rhs, "operands of '=' should not be empty");
        Indent(dep), std::cout << ARGUMENT "LHS:\n";
        lhs->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "RHS:\n";
        rhs->Print(dep + 1);
    }

    void CallExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Call\n";
        PrintInfo(dep);
        Utils::Assert(func && args,
                      "function ot arguments should not be empty");
        Indent(dep), std::cout << ARGUMENT "Function:\n";
        func->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "Arguments:\n";
        args->Print(dep + 1);
    }

    void ArrayExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Array\n";
        PrintInfo(dep);
        Utils::Assert(lhs && rhs, "operands of '[]' should not be empty");
        Indent(dep), std::cout << ARGUMENT "LHS:\n";
        lhs->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "RHS:\n";
        rhs->Print(dep + 1);
    }

    void OperatorExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Operator: ";
        std::cout << VALUE << (char)op << "\n";
        PrintInfo(dep);
        Utils::Assert(lhs && rhs, "operands of %c should not be empty",
                      (char)op);
        Indent(dep), std::cout << ARGUMENT "LHS:\n";
        lhs->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "RHS:\n";
        rhs->Print(dep + 1);
    }

    void TernaryExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Ternary\n";
        PrintInfo(dep);
        Utils::Assert(cond && tVal && fVal,
                      "operands of '?:' should not be empty");
        Indent(dep), std::cout << ARGUMENT "Condition:\n";
        cond->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "True:\n";
        tVal->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "False:\n";
        fVal->Print(dep + 1);
    }

    void LiteralExpr::Print(int dep) const {
        Indent(dep);
        if (value.type == Type::Basic::Int) {
            std::cout << NAME "Integer Literal: ";
            std::cout << VALUE << value.v.i << "\n";
        } else if (value.type == Type::Basic::Float) {
            std::cout << NAME "Float Literal: ";
            std::cout << VALUE << value.v.f << "\n";
        }
        PrintInfo(dep);
    }

    void VoidExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Void\n";
        PrintInfo(dep);
    }

    void IdentifierExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Identifier: ";
        std::cout << VALUE << name << "\n";
        PrintInfo(dep);
    }

    void InitializerExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Initializer\n";
        PrintInfo(dep);
        for (auto item : items) item->Print(dep + 1);
        for (auto [pos, item] : parsed) {
            Indent(dep), std::cout << ARGUMENT << "Offset " << pos << ":\n";
            item->Print(dep + 1);
        }
    }
    void CastExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "Cast: ";
        std::cout << VALUE << (char)from << " -> " << (char)to << "\n";
        PrintInfo(dep);
        origin->Print(dep + 1);
    }

    void AddrExpr::Print(int dep) const {
        Indent(dep), std::cout << NAME "AddrExpr\n";
        PrintInfo(dep);
        Utils::Assert(base && offset,
                      "operands of AddrExpr should not be empty");
        Indent(dep), std::cout << ARGUMENT "Base:\n";
        base->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "Offset:\n";
        offset->Print(dep + 1);
    }

    void ArgList::Print(int dep) const {
        Indent(dep), std::cout << NAME "ArgList\n";
        for (auto item : items) item->Print(dep + 1);
    }

    void Block::Print(int dep) const {
        Indent(dep), std::cout << NAME "Block\n";
        if (!local.empty()) {
            Indent(dep), std::cout << ARGUMENT "Local:\n";
            for (auto sym : local) {
                Indent(dep + 1), std::cout << VALUE << sym->name << DEF ": ";
                std::cout << VALUE << sym->allocMem << DEF << ", ";
                sym->info->Print(0);
            }
        }
        Indent(dep), std::cout << ARGUMENT "Content:\n";
        for (auto item : items) item->Print(dep + 1);
    }

    void VarDecl::Print(int dep) const {
        Indent(dep), std::cout << NAME "VarDecl: ";
        std::cout << VALUE << (isConst ? "const " : "");
        std::cout << (char)baseType;
        std::cout << ' ' << name << "\n";
        if (!shape.empty()) {
            Indent(dep), std::cout << ARGUMENT "Shape:\n";
            for (auto x : shape) x->Print(dep + 1);
        }
        if (init) {
            Indent(dep), std::cout << ARGUMENT "Init:\n";
            init->Print(dep + 1);
        }
    }

    void FuncDecl::Print(int dep) const {
        Indent(dep), std::cout << NAME "FuncDecl: " VALUE;
        std::cout << (char)retType;
        std::cout << ' ' << name << "\n";
        if (!args.empty()) {
            Indent(dep), std::cout << ARGUMENT "Arguments:\n";
            for (auto arg : args) arg->Print(dep + 1);
        }
        Indent(dep), std::cout << ARGUMENT "Content:\n";
        content->Print(dep + 1);
    }

    void IfStmt::Print(int dep) const {
        Indent(dep), std::cout << NAME "If\n";
        Utils::Assert(cond && tVal && fVal,
                      "if condition and content should not be empty");
        Indent(dep), std::cout << ARGUMENT "Condition:\n";
        cond->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "True:\n";
        tVal->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "False:\n";
        fVal->Print(dep + 1);
    }

    void WhileStmt::Print(int dep) const {
        Indent(dep), std::cout << NAME "While\n";
        Utils::Assert(cond && content,
                      "while condition and content should not be empty");
        Indent(dep), std::cout << ARGUMENT "Condition:\n";
        cond->Print(dep + 1);
        Indent(dep), std::cout << ARGUMENT "Content:\n";
        content->Print(dep + 1);
    }

    void ContinueStmt::Print(int dep) const {
        Indent(dep), std::cout << NAME "Continue\n";
    }

    void BreakStmt::Print(int dep) const {
        Indent(dep), std::cout << NAME "Break\n";
    }

    void ReturnStmt::Print(int dep) const {
        Indent(dep), std::cout << NAME "Return\n";
        if (ret) ret->Print(dep + 1);
    }

}  // namespace AST
